import { BaiduModelType } from "./enum";

/**
 * 定义 baidu 的模型配置
 */
export interface BaiduConfig {
    /**
     * 较高的数值会使输出更加随机，而较低的数值会使其更加集中和确定，默认0.8，范围 (0, 1.0]
     */
    temperature?: number;
    /**
     * 影响输出文本的多样性，取值越大，生成文本的多样性越强，默认0.8，取值范围 [0, 1.0]
     */
    top_p?: number;
    /**
     * 通过对已生成的token增加惩罚，减少重复生成的现象，默认1.0，取值范围：[1.0, 2.0]
     */
    penalty_score?: number;
    /**
     * 是否以流式接口的形式返回数据，默认false
     */
    stream?: boolean;
    /**
     * 是否开启系统记忆，默认false
     */
    enable_system_memory?: boolean;
    /**
     * 系统记忆ID，用于读取对应ID下的系统记忆
     */
    system_memory_id?: string;
    /**
     * 模型人设，主要用于人设设定
     */
    system?: string;
    /**
     * 生成停止标识，当模型生成结果以stop中某个元素结尾时，停止文本生成
     */
    stop?: string[];
    /**
     *  是否强制关闭实时搜索功能，默认false
     */
    disable_search?: boolean;
    /**
     * 是否开启上角标返回，默认false
     */
    enable_citation?: boolean;
    /**
     * 是否开启搜索溯源功能，默认false
     */
    enable_trace?: boolean;
    /**
     * 指定模型最大输出token数，范围[2, 2048]，默认1024
     */
    max_output_tokens?: number;
    /**
     * 指定响应内容的格式，默认text
     */
    response_format?: "json_object" | "text";
    /**
     * 用户IP
     */
    user_ip?: string;
    /**
     * 表示最终用户的唯一标识符
     */
    user_id?: string;
}

/**
 * 默认的 ChatGPTConfig 配置
 */
export const defaultBaiduConfig: BaiduConfig = {
    temperature: 0.8,
    top_p: 1.0,
    stream: false,
    stop: undefined,
    response_format: "text"
};

/**
 * 各个模型的 token 数量限制
 */
export const modelTokenLimit: Record<string, number> = {
    [BaiduModelType.ERNIE_4_0_8K]: 3000,
    [BaiduModelType.ERNIE_4_0_TURBO_8K]: 3000,
    [BaiduModelType.ERNIE_SPEED_PRO_128K]: 128000
};
